package com.ruoyi.framework.websocket;

import com.ruoyi.common.utils.DateUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * 高级自定义 WebSocket 处理器
 * 支持用户关联、消息广播、定向发送和连接状态管理
 */
@Slf4j
@Component
public class AdvancedWebSocketHandler extends TextWebSocketHandler {

    // 存储所有活跃连接（线程安全集合）
    private final Set<WebSocketSession> activeSessions = new CopyOnWriteArraySet<>();
    
    // 用户ID与WebSocket会话的映射（支持定向消息）
    private final Map<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();
    
    // 会话ID与用户ID的反向映射（便于连接关闭时清理）
    private final Map<String, String> sessionUserMap = new ConcurrentHashMap<>();

    /**
     * 连接建立时触发
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 添加会话到活跃连接集合
        activeSessions.add(session);
        String sessionId = session.getId();
        logConnectionEvent("建立连接", sessionId);

        // 从会话属性获取用户ID（需在握手拦截器中设置）
        String userId = (String) session.getAttributes().get("userId");
        if (userId != null && !userId.isEmpty()) {
            // 存储用户与会话的双向映射
            userSessionMap.put(userId, session);
            sessionUserMap.put(sessionId, userId);
            
            // 向客户端发送欢迎消息
            sendMessageToSession(session, "系统通知", "欢迎回来，用户[" + userId + "]！当前在线人数: " + activeSessions.size());
        } else {
            // 匿名用户处理
            sendMessageToSession(session, "系统通知", "连接成功！sessionId: " + sessionId + "，当前在线人数: " + activeSessions.size());
        }
    }

    /**
     * 处理客户端发送的文本消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String sessionId = session.getId();
        String payload = message.getPayload();
        String userId = sessionUserMap.getOrDefault(sessionId, "匿名用户");

        System.out.printf("[%s] 收到用户[%s]的消息: %s%n", 
                LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_TIME), 
                userId, payload);

        // 1. 回复消息给发送者
        sendMessageToSession(session, "服务器回复", "已发送");

        // 2. 广播消息给所有在线用户（排除发送者）
        broadcastMessage(userId, payload, session);
    }

    /**
     * 连接关闭时触发
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String sessionId = session.getId();
        // 从活跃连接中移除
        activeSessions.remove(session);
        
        // 清理用户映射
        String userId = sessionUserMap.get(sessionId);
        if (userId != null) {
            userSessionMap.remove(userId);
            sessionUserMap.remove(sessionId);
        }

        logConnectionEvent("关闭连接", sessionId + "，原因: " + status.getReason());
        // 广播用户离开消息
        broadcastSystemMessage("用户[" + (userId != null ? userId : "匿名用户") + "]已离开，当前在线人数: " + activeSessions.size());
    }

    /**
     * 处理传输错误
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String sessionId = session.getId();
        String errorMsg = "连接错误: " + exception.getMessage();
        System.err.println("[" + sessionId + "] " + errorMsg);
        
        // 尝试向客户端发送错误通知
        if (session.isOpen()) {
            sendMessageToSession(session, "系统错误", errorMsg);
        }
    }

    /**
     * 向指定会话发送消息
     */
    private void sendMessageToSession(WebSocketSession session, String type, String content) throws IOException {
        if (session.isOpen()) {
            // 构建结构化消息
            String message = String.format("{\"type\":\"%s\",\"time\":\"%s\",\"content\":\"%s\"}",
                    type,
                    DateUtils.getLocalTime(),
                    content);
            session.sendMessage(new TextMessage(message));
        }
    }

    /**
     * 广播消息给所有在线用户（可排除发送者）
     */
    public void broadcastMessage(String sender, String content, WebSocketSession excludeSession) throws IOException {
        String message = String.format("{\"type\":\"广播消息\",\"sender\":\"%s\",\"time\":\"%s\",\"content\":\"%s\"}",
                sender,
                DateUtils.getLocalTime(),
                content);

        for (WebSocketSession session : activeSessions) {
            if (session.isOpen() && (excludeSession == null || !session.getId().equals(excludeSession.getId()))) {
                session.sendMessage(new TextMessage(message));
            }
        }
    }

    /**
     * 向特定用户发送定向消息
     */
    public boolean sendToUser(String userId, String content) throws IOException {
        WebSocketSession session = userSessionMap.get(userId);
        if (session != null && session.isOpen()) {
            sendMessageToSession(session, "私人消息", content);
            return true;
        }
        return false;
    }

    /**
     * 广播系统消息
     */
    public void broadcastSystemMessage(String content) throws IOException {
        broadcastMessage("系统", content, null);
    }

    /**
     * 记录连接事件日志
     */
    private void logConnectionEvent(String eventType, String details) {
        System.out.printf("[%s] WebSocket %s - %s，当前在线人数: %d%n",
                LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_TIME),
                eventType,
                details,
                activeSessions.size());
    }

    // 辅助方法：获取当前在线人数
    public int getOnlineCount() {
        return activeSessions.size();
    }

    // 辅助方法：检查用户是否在线
    public boolean isUserOnline(String userId) {
        WebSocketSession session = userSessionMap.get(userId);
        return session != null && session.isOpen();
    }
}
